import onnx
from onnx import helper

def helper():
    model = onnx.load('yolov11_xw_tiny_100_best.onnx')
    # Create a new graph with modified nodes
    new_nodes = []
    for node in model.graph.node:
        if node.op_type == "Reshape":
            # Replace with a compatible version (e.g., opset 13)
            new_node = helper.make_node(
                "Reshape",
                inputs=node.input,
                outputs=node.output,
                name=node.name
            )
            new_nodes.append(new_node)
        else:
            new_nodes.append(node)
    
    # Rebuild the model
    new_model = helper.make_model(
        helper.make_graph(
            new_nodes,
            model.graph.name,
            model.graph.input,
            model.graph.output,
        ),
        opset_imports=[helper.make_opsetid("", 17)]  # Downgrade opset
    )
    
    onnx.save(new_model, "yolov11_xw_tiny_100_best_opset17_fixed.onnx")

def change_op_version():
    model = onnx.load('yolov11_xw_tiny_100_best-cut.onnx')
    model.opset_import[0].version = 17
    model.ir_version = 10
    onnx.save_model(model, 'yolov11_xw_tiny_100_best-cut_op17.onnx')

def get_cut_out():
    input_path = 'yolov11_xw_tiny_100_best.onnx'
    output_path = 'yolov11_xw_tiny_100_best-cut.onnx'
    input_names = ['images']
    output_names = ["/model.23/Concat_output_0", "/model.23/Concat_1_output_0", "/model.23/Concat_2_output_0"]
    onnx.utils.extract_model(input_path, output_path, input_names, output_names)

#get_cut_out()
change_op_version()
